#include <CL/cl.h>
#include <iostream>
#include <vector>
#include <fstream>
#include <cmath>
#include <cassert>

std::string loadKernelSource(const std::string& filename) {
    std::string kernel_path = KERNELS_DIR "/" + filename;
    std::ifstream file(kernel_path);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to open kernel file: " + kernel_path);
    }
    return std::string((std::istreambuf_iterator<char>(file)), 
                      std::istreambuf_iterator<char>());
}

int main() {
    cl_int err;
    cl_uint num_platforms;
    cl_platform_id platform = nullptr;
    cl_device_id device = nullptr;

    // --- 1. 获取平台与设备 ---
    err = clGetPlatformIDs(1, &platform, &num_platforms);
    assert(err == CL_SUCCESS);

    err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, 1, &device, nullptr);
    if (err != CL_SUCCESS) {
        std::cerr << "⚠️ 没检测到 GPU，改用 CPU...\n";
        err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_CPU, 1, &device, nullptr);
        assert(err == CL_SUCCESS);
    }

    // --- 2. 创建上下文与命令队列 ---
    cl_context context = clCreateContext(nullptr, 1, &device, nullptr, nullptr, &err);
    assert(err == CL_SUCCESS);

    cl_command_queue queue = clCreateCommandQueue(context, device, 0, &err);
    assert(err == CL_SUCCESS);

    // --- 3. 读取并构建内核 ---
    std::string source = loadKernelSource("swiglu.cl");
    const char *src = source.c_str();
    size_t length = source.size();
    cl_program program = clCreateProgramWithSource(context, 1, &src, &length, &err);
    assert(err == CL_SUCCESS);

    err = clBuildProgram(program, 1, &device, nullptr, nullptr, nullptr);
    if (err != CL_SUCCESS) {
        size_t log_size;
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, nullptr, &log_size);
        std::vector<char> log(log_size);
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, log_size, log.data(), nullptr);
        std::cerr << "❌ 内核编译失败:\n" << log.data() << std::endl;
        return 1;
    }

    cl_kernel kernel = clCreateKernel(program, "kernel_swiglu", &err);
    assert(err == CL_SUCCESS);

    // --- 4. 准备输入输出数据 ---
    const int N = 16;
    std::vector<float> src0(N), src1(N), dst(N, 0.0f);

    for (int i = 0; i < N; ++i) {
        src0[i] = 0.1f * i - 0.8f;
        src1[i] = 1.0f + 0.05f * i;
    }

    cl_mem buf_src0 = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*N, src0.data(), &err);
    cl_mem buf_src1 = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*N, src1.data(), &err);
    cl_mem buf_dst  = clCreateBuffer(context, CL_MEM_WRITE_ONLY, sizeof(float)*N, nullptr, &err);
    assert(err == CL_SUCCESS);

    // --- 5. 设置参数 ---
    size_t offset = 0;
    ulong nb = sizeof(float) * N;

    int ne0 = N;
    int ne00_off = 0, ne10_off = 0;

    int arg = 0;
    err  = clSetKernelArg(kernel, arg++, sizeof(cl_mem), &buf_src0);
    err |= clSetKernelArg(kernel, arg++, sizeof(ulong), &offset);
    err |= clSetKernelArg(kernel, arg++, sizeof(cl_mem), &buf_src1);
    err |= clSetKernelArg(kernel, arg++, sizeof(ulong), &offset);
    err |= clSetKernelArg(kernel, arg++, sizeof(cl_mem), &buf_dst);
    err |= clSetKernelArg(kernel, arg++, sizeof(ulong), &offset);
    err |= clSetKernelArg(kernel, arg++, sizeof(ulong), &nb);
    err |= clSetKernelArg(kernel, arg++, sizeof(ulong), &nb);
    err |= clSetKernelArg(kernel, arg++, sizeof(int), &ne0);
    err |= clSetKernelArg(kernel, arg++, sizeof(ulong), &nb);
    err |= clSetKernelArg(kernel, arg++, sizeof(int), &ne00_off);
    err |= clSetKernelArg(kernel, arg++, sizeof(int), &ne10_off);
    assert(err == CL_SUCCESS);

    // --- 6. 启动 kernel ---
    size_t global_work_size[1] = { (size_t)N };
    size_t local_work_size[1]  = { 8 }; // 可调整
    err = clEnqueueNDRangeKernel(queue, kernel, 1, nullptr, global_work_size, local_work_size, 0, nullptr, nullptr);
    assert(err == CL_SUCCESS);

    clFinish(queue);

    // --- 7. 读取输出 ---
    err = clEnqueueReadBuffer(queue, buf_dst, CL_TRUE, 0, sizeof(float)*N, dst.data(), 0, nullptr, nullptr);
    assert(err == CL_SUCCESS);

    // --- 8. 打印结果 ---
    std::cout << "=== SWIGLU 输出 ===\n";
    for (int i = 0; i < N; ++i) {
        float silu = src0[i] / (1.0f + std::exp(-src0[i]));
        float expected = silu * src1[i];
        std::cout << i << ": GPU=" << dst[i] << "  CPU=" << expected << std::endl;
    }

    // --- 9. 清理 ---
    clReleaseMemObject(buf_src0);
    clReleaseMemObject(buf_src1);
    clReleaseMemObject(buf_dst);
    clReleaseKernel(kernel);
    clReleaseProgram(program);
    clReleaseCommandQueue(queue);
    clReleaseContext(context);

    std::cout << "✅ 测试完成！" << std::endl;
    return 0;
}
